package cn.cnic.security.nssap.stearm

import cn.cnic.security.nssap.stearm.CommonTools.getClass
import com.alibaba.fastjson.{JSON, JSONArray, JSONObject}
import org.apache.flink.configuration.Configuration
import org.apache.flink.streaming.api.functions.sink.{RichSinkFunction, SinkFunction}

import java.io.InputStream
import java.lang
import java.sql.{Connection, DriverManager, PreparedStatement, SQLType, Statement, Timestamp, Types}
import java.util.Properties

/**
  * Mysql 8.0
  */
class MySQLJdbcSink extends RichSinkFunction[(String,String)]{

  // 定义sql连接、预编译器
  var conn: Connection = _
  var alarmStmt: PreparedStatement = _
//  var sandboxStmt: PreparedStatement = _
//  var sandboxMailStmt: PreparedStatement = _
//  var yunsuoStmt: PreparedStatement = _

  var num :Int = _

  val batchSize:Int = 1000

  override def open(parameters: Configuration): Unit = {

    //从配置文件读取kafka设置
    val properties = new Properties()
    val jdbcInput: InputStream = getClass.getClassLoader.getResourceAsStream("jdbc.properties")
    properties.load(jdbcInput)

    Class.forName(properties.getProperty("driverName"))
    val url: String = properties.getProperty("url")
    val username: String = properties.getProperty("username")
    val password: String = properties.getProperty("password")
    conn = DriverManager.getConnection(url, username, password)

    alarmStmt = conn.prepareStatement("INSERT INTO `nssap`.`qxn_warn_log_alarm`(`first_access_time`, `hazard_level`, `type_chain`, `attck`, `repeat_count`, `uri`, `x_forwarded_for`, `rsp_status`, `alarm_source`, `alarm_sip`, `threat_name`, `group_id`, `attack_org`, `attack_sip`, `host`, `host_state`, `ioc`, `access_time`, `attack_chain` ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)")
//    sandboxStmt = conn.prepareStatement("")
//    sandboxMailStmt = conn.prepareStatement("")
//    yunsuoStmt = conn.prepareStatement("")
    //禁止自动提交
    conn.setAutoCommit(false) //
    num = 0
  }

  override def invoke(value: (String, String), context: SinkFunction.Context): Unit = {

    value._1 match {
      case "alarm" => {
        val map: JSONObject = JSON.parseObject(value._2)

        val first_access_time: lang.Long = map.getLong("first_access_time")
        alarmStmt.setTimestamp(1,new Timestamp(first_access_time))

        val hazard_level = map.getByte("hazard_level")
        alarmStmt.setByte(2,hazard_level)

        val type_chain: String = map.getString("type_chain")
        alarmStmt.setString(3,type_chain)

        val attck: String = map.getString("attck")
        alarmStmt.setString(4,attck)

        val repeat_count: Integer = map.getInteger("repeat_count")
        alarmStmt.setInt(5,repeat_count)

        val uri: String = map.getString("uri")
        alarmStmt.setString(6,uri)

        val x_forwarded_for: String = map.getString("x_forwarded_for")
        if(x_forwarded_for !=null && x_forwarded_for.length > 0){
          alarmStmt.setString(7,x_forwarded_for)
        }else{
          alarmStmt.setNull(7,Types.VARCHAR)
        }

        val rsp_status: JSONArray = map.getJSONArray("rsp_status")
        //数组，用逗号分隔保存到一列
        if(rsp_status !=null && rsp_status.size() > 0){
          val sb = new StringBuilder
          for(i <- 0 until rsp_status.size() ) {
            sb.append(rsp_status.getString(i))
            if(rsp_status.size() - i != 1) {
              sb.append(",")
            }
          }
          alarmStmt.setString(8,sb.toString())
        }else{
          alarmStmt.setNull(8,Types.VARCHAR)
        }

        val alarm_source: Integer = map.getInteger("alarm_source")
        if(alarm_source != null){
          alarmStmt.setInt(9,alarm_source)
        }else{
          alarmStmt.setNull(9,Types.INTEGER)
        }

        val alarm_sip: String = map.getString("alarm_sip")
        alarmStmt.setString(10,alarm_sip)

        val threat_name: String = map.getString("threat_name")
        alarmStmt.setString(11,threat_name)

        val group_id: Integer = map.getInteger("group_id")
        if(group_id != null){
          alarmStmt.setInt(12,group_id)
        }else{
          alarmStmt.setNull(12,Types.INTEGER)
        }

        val attack_org: String = map.getString("attack_org")
        if(attack_org != null && attack_org.length > 0){
          alarmStmt.setString(13,attack_org)
        }else{
          alarmStmt.setNull(13,Types.VARCHAR)
        }

        val attack_sip: String = map.getString("attack_sip")
        alarmStmt.setString(14,attack_sip)

        val host: String = map.getString("host")
        alarmStmt.setString(15,host)

        val host_state: Integer = map.getInteger("host_state")
        if(host_state != null){
          alarmStmt.setInt(16,host_state)
        }else{
          alarmStmt.setNull(16,Types.INTEGER)
        }

        val ioc: String = map.getString("ioc")
        alarmStmt.setString(17,ioc)

        val access_time: lang.Long = map.getLong("access_time")
        alarmStmt.setTimestamp(18,new Timestamp(access_time))

        val attack_chain: String = map.getString("attack_chain")
        alarmStmt.setString(19,attack_chain)

        alarmStmt.addBatch()
      }
      case "sandbox" =>{

      }
      case "sandbox_mail" =>{

      }
      case "yunsuo" =>{

      }
    }
    num += 1
    //执行一次提交
    if(batchSize % 10 == 0){
      val ints: Array[Int] = alarmStmt.executeBatch()
//      println("循环 插入结果 "+ints.mkString(","))
      alarmStmt.clearBatch()
      conn.commit()
    }

  }



  override def close(): Unit = {
    val ints: Array[Int] = alarmStmt.executeBatch()
//    println("close插入结果 "+ints.mkString(","))
    conn.commit() //提交事务

    //关闭资源
    if(alarmStmt != null)
      alarmStmt.close()

    if(conn != null)
      conn.close()
  }
}
